import torch
import torch.nn as nn


# 定义卷积块，保持不变
class ConvBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(ConvBlock, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
        )

    def forward(self, x):
        return self.conv(x)


# 定义 MEnet 中可能有的特殊模块，调整初始化参数以更好适配网络结构
class MENetModule(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(MENetModule, self).__init__()
        # 多尺度卷积与注意力机制示例
        self.conv_branch1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
        self.conv_branch2 = nn.Conv2d(in_channels, out_channels, kernel_size=5, padding=2)
        self.bn = nn.BatchNorm2d(2 * out_channels)
        self.relu = nn.ReLU(inplace=True)
        self.reduce_conv = nn.Conv2d(2 * out_channels, out_channels, kernel_size=1)

    def forward(self, x):
        out1 = self.conv_branch1(x)
        out2 = self.conv_branch2(x)
        out = torch.cat([out1, out2], dim=1)
        out = self.bn(out)
        out = self.relu(out)
        out = self.reduce_conv(out)
        return out


# 定义完整的 MEnet-VNet
class MENetVNet(nn.Module):
    def __init__(self, in_channels=1, out_channels=1):
        super(MENetVNet, self).__init__()
        # 编码器部分，每个编码器后接 MEnet 模块
        self.encoder1 = nn.Sequential(
            ConvBlock(in_channels, 64),
            MENetModule(64, 64)
        )
        self.encoder2 = nn.Sequential(
            ConvBlock(64, 128),
            MENetModule(128, 128)
        )
        self.encoder3 = nn.Sequential(
            ConvBlock(128, 256),
            MENetModule(256, 256)
        )
        self.encoder4 = nn.Sequential(
            ConvBlock(256, 512),
            MENetModule(512, 512)
        )

        self.pool = nn.MaxPool2d(2, 2)
        # 瓶颈层，包含卷积块与 MEnet 模块
        self.bottleneck = nn.Sequential(
            ConvBlock(512, 1024),
            MENetModule(1024, 1024)
        )

        # 上采样与解码器部分，包含转置卷积、卷积块与 MEnet 模块
        self.upconv4 = nn.ConvTranspose2d(1024, 512, kernel_size=2, stride=2)
        self.decoder4 = nn.Sequential(
            ConvBlock(1024, 512),
            MENetModule(512, 512)
        )
        self.upconv3 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)
        self.decoder3 = nn.Sequential(
            ConvBlock(512, 256),
            MENetModule(256, 256)
        )
        self.upconv2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)
        self.decoder2 = nn.Sequential(
            ConvBlock(256, 128),
            MENetModule(128, 128)
        )
        self.upconv1 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
        self.decoder1 = nn.Sequential(
            ConvBlock(128, 64),
            MENetModule(64, 64)
        )

        self.out_conv = nn.Conv2d(64, out_channels, kernel_size=1)

    def forward(self, x):
        # 编码器前向传播
        enc1 = self.encoder1(x)
        enc2 = self.encoder2(self.pool(enc1))
        enc3 = self.encoder3(self.pool(enc2))
        enc4 = self.encoder4(self.pool(enc3))

        # 瓶颈层前向传播
        bottleneck = self.bottleneck(self.pool(enc4))

        # 解码器前向传播，融合跳跃连接特征
        dec4 = self.decoder4(torch.cat([self.upconv4(bottleneck), enc4], dim=1))
        dec3 = self.decoder3(torch.cat([self.upconv3(dec4), enc3], dim=1))
        dec2 = self.decoder2(torch.cat([self.upconv2(dec3), enc2], dim=1))
        dec1 = self.decoder1(torch.cat([self.upconv1(dec2), enc1], dim=1))

        return torch.sigmoid(self.out_conv(dec1))